package com.ruoyi.ade.service.impl;

import com.alibaba.fastjson2.JSONArray;
import com.alibaba.fastjson2.JSONObject;
import com.ruoyi.ade.domain.*;
import com.ruoyi.ade.domain.AdeKnowledge;
import com.ruoyi.ade.llm.agent.DashScopeAgent;
import com.ruoyi.ade.llm.agent.OpenAIAgent;
import com.ruoyi.ade.llm.agent.QianfanAgent;
import com.ruoyi.ade.llm.model.*;
import com.ruoyi.ade.llm.service.IAdeKnowledgeEmbedService;
import com.ruoyi.ade.llm.service.IAdeService;
import com.ruoyi.ade.mapper.*;
import com.ruoyi.ade.model.AdeKnowledgeRetriveRequest;
import com.ruoyi.ade.model.AdeRecognitionRequest;
import com.ruoyi.ade.service.IAdeKnowledgeCacheService;
import com.ruoyi.ade.service.IAdeLlmExecuteService;
import com.ruoyi.ade.tool.AdeCompareUtils;
import com.ruoyi.common.utils.StringUtils;
import lombok.extern.slf4j.Slf4j;
import org.checkerframework.checker.units.qual.A;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

@Service
@Slf4j
public class AdeLlmExecuteServiceImpl implements IAdeLlmExecuteService {

    @Autowired
    private AdeLlmLabMapper adeLlmLabMapper;
    @Autowired
    private AdeLlmLabResultMapper adeLlmLabResultMapper;
    @Autowired
    private AdePromptMapper adePromptMapper;
    @Autowired
    private StandardMapper standardMapper;
    @Autowired
    private AdeKnowledgeBaseMapper adeKnowledgeBaseMapper;
    @Autowired
    private IAdeService adeService;

    @Autowired
    private IAdeKnowledgeEmbedService adeKnowledgeEmbedService;

    @Autowired
    private IAdeKnowledgeCacheService adeKnowledgeCacheService;


    @Override
    public void execute(Long id) {

        AdeLlmLab adeLlmLab = adeLlmLabMapper.selectAdeLlmLabById(id);
        //检查ID是否存在
        if (adeLlmLab == null) throw new RuntimeException("No lab found for:" + id);

        //失败的未启动的可以执行
        if (!AdeLab.STATUS_READY.equals(adeLlmLab.getStatus()) && !AdeLab.STATUS_FAILED.equals(adeLlmLab.getStatus()))
            throw new RuntimeException("Current lab status exception:" + adeLlmLab.getStatus());

        //检查prompt是否存在
        AdePrompt adePrompt = adePromptMapper.selectAdePromptById(adeLlmLab.getPromptId());
        if (adePrompt == null) throw new RuntimeException("No prompt found for:" + adeLlmLab.getPromptId());

        //更新状态为运行中
        updateLlmLab(id, AdeLab.STATUS_RUNNING,null,null);
        log.info("LLMLab id:{}, Started update status:running", id);

        //开始执行任务 线程
        new Thread(new Runnable() {

            @Override
            public void run() {
                int finished = 0;  //统计完成数量
                try{

                    //获取符合条件的 需要分析的标准数据集
                    List<Standard> avis = getAviStandard(adeLlmLab);
                    log.info("LLMLab id:{}, Get avi standards size:{}", id, avis.size());

                    //获取已经完成的记录
                    AdeLlmLabResult cond = new AdeLlmLabResult();
                    cond.setLlmTaskId(id);
                    List<AdeLlmLabResult> results = adeLlmLabResultMapper.selectAdeLlmLabResultList(cond);
                    finished = results.size();
                    log.info("LLMLab id:{}, Get finished results size:{}", id, finished);
                    Map<Long,Byte> cache = new HashMap<>();
                    results.forEach(r->cache.put(r.getStandardId(),(byte)01));

                    //更新任务需要完成的数量
                    updateLlmLab(id,null,(long)finished,(long)avis.size());
                    log.info("LLMLab id:{}, Update total size:{}", id, avis.size());

                    //获取RAG检索配置
                    AdeKnowledgeRetriveRequest ragRequest = JSONObject.parseObject(adeLlmLab.getRagConfig(), AdeKnowledgeRetriveRequest.class);
                    log.info("LLMLab id:{}, RAG request:{}", id, JSONObject.toJSONString(ragRequest));

                    String userPrompt, systemPrompt;
                    List<AdeLlmLabResult> adeLlmLabResults = new ArrayList<>();
                    //逐个分析数据集
                    for(Standard standard : avis){

                        if(cache.get(standard.getId())!=null) continue;  //已经分析过了

                        log.info("LLMLab id:{}, Start analyze standard:{}", id, standard.getId());
                        List<Ade> references = JSONArray.parseArray(standard.getAdes(),Ade.class);

                        //记录开始时间
                        long cur = System.currentTimeMillis();
                        //构建分析结果
                        AdeLlmLabResult adeLlmLabResult = new AdeLlmLabResult();
                        adeLlmLabResult.setStandardId(standard.getId());
                        adeLlmLabResult.setLlmTaskId(adeLlmLab.getId());
                        String remark = "";
                        try {
                            float avgSimScore = 0f;
                            int knowledgeCount = 0;
                            //构建用户提示
                            userPrompt = String.format(adePrompt.getUserPrompt(), standard.getContent());
                            systemPrompt = adePrompt.getSystemPrompt();

                            //RAG 模式下需要检索知识  构建系统提示
                            if (AdePrompt.PROMPT_TYPE_RAG.equals(adeLlmLab.getPromptType())) {

                                //RAG知识检索，基于缓存策略
                                List<AdeRetrivedKnowledge> knowledges = adeKnowledgeCacheService.retrieveKnowledge(standard.getId(), standard.getContent(), ragRequest);
                                for (AdeRetrivedKnowledge k : knowledges) avgSimScore += k.getSimScore();
                                knowledgeCount = knowledges.size();
                                if(knowledgeCount != 0)
                                     avgSimScore = avgSimScore / knowledgeCount;
                                //构建系统提示
                                systemPrompt = String.format(adePrompt.getSystemPrompt(), AdeRetrivedKnowledge.formatKnowledge(knowledges));
                            }
                            adeLlmLabResult.setAvgSimscore(avgSimScore);
                            adeLlmLabResult.setKnowledgeCount(knowledgeCount);

                            //算法识别ADE结果
                            List<Ade> ades = adeService.recognize(adeLlmLab.getLlm(),systemPrompt,userPrompt,true);
                            //跟标准结果进行比对分析  获取匹配等级
                            String level = AdeCompareUtils.compare(references, ades);
                            adeLlmLabResult.setMatchLevel(level);
                            adeLlmLabResult.setAdes(JSONArray.toJSONString(ades));
                            adeLlmLabResult.setStatus(AdeLlmLabResult.STATUS_OK);

                        }catch (Exception e){
                            remark = e.getMessage();
                            adeLlmLabResult.setAdes("[]");
                            adeLlmLabResult.setMatchLevel(AdeCompareUtils.LEVEL_LN);
                            adeLlmLabResult.setStatus(AdeLlmLabResult.STATUS_ERR);
                            log.error("LLMLab id:{}, Exception:{}", id, e.getMessage());
                        }
                        adeLlmLabResult.setRemark(remark);
                        adeLlmLabResult.setTimeConsumed((System.currentTimeMillis() - cur)/1000); //转换成秒
                        adeLlmLabResults.add(adeLlmLabResult);
                        finished++;

                        if(adeLlmLabResults.size()>=5){
                            //批量插入
                            adeLlmLabResultMapper.batchInsertAdeLlmLabResult(adeLlmLabResults);
                            adeLlmLabResults.clear();
                            //更新任务进度
                            updateLlmLab(id, null, (long)finished, null);
                            log.info("LLMLab id:{}, Finished update progress:{}", id, finished);
                        }

                        //更新标准数据集
                        updateStandardAdes(standard.getId(), adeLlmLabResult.getAdes(), adeLlmLabResult.getMatchLevel(), adeLlmLab.getLlm(), adeLlmLab.getPromptType());
                    }

                    //最后检查 并 插入
                    if(adeLlmLabResults.size()>0){
                        adeLlmLabResultMapper.batchInsertAdeLlmLabResult(adeLlmLabResults);
                        adeLlmLabResults.clear();
                    }
                }catch (Exception e){
                    updateLlmLab(id, AdeLab.STATUS_FAILED, null, null);
                    log.error("LLMLab id:{}, Exception:{}", id, e.getMessage());
                    return;
                }

                //执行完成
                updateLlmLab(id, AdeLab.STATUS_SUCCESS, (long)finished, null);
                log.info("LLMLab id:{}, Finished update status:success", id);
            }
        }).start();

    }

    private void updateStandardAdes(Long standardId, String ades, String matchLevel, String llm, String promptType){
        Standard update = new Standard();
        update.setId(standardId);
        switch (llm){
            case QianfanAgent.model_deepseek_v3:
                if(AdePrompt.PROMPT_TYPE_NAG.equals(promptType)) {
                    update.setNagAdesDeepseek(ades);
                    update.setNagAdesDeepseekRes(matchLevel);
                }
                else if(AdePrompt.PROMPT_TYPE_SAG.equals(promptType)){
                    update.setAdesDeepseek(ades);
                    update.setAdesDeepseekRes(matchLevel);
                }
                else if(AdePrompt.PROMPT_TYPE_RAG.equals(promptType)){
                    update.setRagAdesDeepseek(ades);
                    update.setRagAdesDeepseekRes(matchLevel);
                }else{
                    log.error("===============Unknown prompt type:{}",promptType);
                }
                break;
            case QianfanAgent.model_ernie35_8k:
                if(AdePrompt.PROMPT_TYPE_NAG.equals(promptType)) {
                    update.setNagAdesEnire(ades);
                    update.setNagAdesEnireRes(matchLevel);
                }
                else if(AdePrompt.PROMPT_TYPE_SAG.equals(promptType)) {
                    update.setAdesEnire(ades);
                    update.setAdesEnireRes(matchLevel);
                }else if(AdePrompt.PROMPT_TYPE_RAG.equals(promptType)) {
                    update.setRagAdesEnire(ades);
                    update.setRagAdesEnireRes(matchLevel);
                }else{
                    log.error("===============Unknown prompt type:{}",promptType);
                }
                break;
            case  OpenAIAgent.model_gpt4o:
                if(AdePrompt.PROMPT_TYPE_NAG.equals(promptType)) {
                    update.setNagAdesGpt(ades);
                    update.setNagAdesGptRes(matchLevel);
                }
                else if(AdePrompt.PROMPT_TYPE_SAG.equals(promptType)) {
                    update.setAdesGpt(ades);
                    update.setAdesGptRes(matchLevel);
                }else if(AdePrompt.PROMPT_TYPE_RAG.equals(promptType)){
                    update.setRagAdesGpt(ades);
                    update.setRagAdesGptRes(matchLevel);
                }else{
                    log.error("===============Unknown prompt type:{}",promptType);
                }
                break;
            case DashScopeAgent.model_qwen_turbo:
                if(AdePrompt.PROMPT_TYPE_NAG.equals(promptType)) {
                    update.setNagAdesQwen(ades);
                    update.setNagAdesQwenRes(matchLevel);
                }
                else if(AdePrompt.PROMPT_TYPE_SAG.equals(promptType)) {
                    update.setAdesQwen(ades);
                    update.setAdesQwenRes(matchLevel);
                }else if(AdePrompt.PROMPT_TYPE_RAG.equals(promptType)){
                    update.setRagAdesQwen(ades);
                    update.setRagAdesQwenRes(matchLevel);
                }else{
                    log.error("===============Unknown prompt type:{}",promptType);
                }
                break;
            default:
                log.error("===============Unknown llm:{}",llm);
                break;
        }
        standardMapper.updateStandard(update);
    }

    private void updateLlmLab(Long id, String status, Long finished, Long total) {
        AdeLlmLab update = new AdeLlmLab();
        update.setId(id);
        if(status != null) update.setStatus(status);
        if(finished != null) update.setFinished(finished);
        if(total != null) update.setTotal(total);
        adeLlmLabMapper.updateAdeLlmLab(update);
    }

    /**
     * @param adeLab
     * @return
     */
    private List<Standard> getAviStandard(AdeLlmLab adeLab) {
        List<Standard> standards = standardMapper.selectStandardList(new Standard());
        log.info("LLMLab Name:{}, Total standards:{}", adeLab.getAlgorithm(), standards.size());
        List<Standard> avis = new ArrayList<>();
        List<String> includes = new ArrayList<>();
        List<String> excludes = new ArrayList<>();

        if (adeLab.getIncludeSources() != null) {
            String[] include = adeLab.getIncludeSources().split("#");
            for (String s : include) {
                if (!StringUtils.isEmpty(s))
                    includes.add(s);
            }
        }
        log.info("LLMLab Name:{}, Include sources:{}", adeLab.getAlgorithm(), JSONArray.toJSONString(includes));

        if (adeLab.getExcludeSources() != null) {
            String[] exclude = adeLab.getExcludeSources().split("#");
            for (String s : exclude) {
                if (!StringUtils.isEmpty(s))
                    excludes.add(s);
            }
        }
        log.info("LLMLab Name:{}, Exclude sources:{}", adeLab.getAlgorithm(), JSONArray.toJSONString(excludes));

        String standardStatus = adeLab.getStandardStatus();

        //分析符合要求的 标准
        for (Standard standard : standards) {
//            System.out.println(standardStatus+"==============:"+standard.getSource()+"\t\t"+!excludes.contains(standard.getSource()));
            if (StringUtils.isEmpty(standard.getSource())) continue; //如果来源是空 则不考虑
            //纳入标准    必须是符合状态
            if (StringUtils.isEmpty(standardStatus) || standard.getStatus().equals(standardStatus)) {
                if (includes.size() == 0 || includes.contains(standard.getSource())) {
                    //排除标准
                    if (excludes.size() == 0 || !excludes.contains(standard.getSource())) {
                        avis.add(standard);
                    }
                }
            }
        }
        return avis;
    }

    /**
     *
     * @param llm                 大模型算法
     * @param systemPrompt        系统提示词
     * @param userPrompt          用户提示词
     * @return
     */
    public List<Ade> recognize(String llm, String systemPrompt, String userPrompt){
        return adeService.recognize(llm, systemPrompt, userPrompt,true);
    }

    @Override
    public List<Ade> recognize(AdeRecognitionRequest request) {
        AdeWithKnowledge r = recognizeWithKnowledge(request);
        return r.getAdes();
    }


    public AdeWithKnowledge recognizeWithKnowledge(AdeRecognitionRequest request){
        AdeWithKnowledge result  = new AdeWithKnowledge();
        //验证EMR是否为空
        if(StringUtils.isEmpty(request.getEmr())) throw new RuntimeException("Clinical Notes(EMR) can not be empty!");
        //验证LLM模型是否为空
        if(StringUtils.isEmpty(request.getLlmModel())) throw new RuntimeException("LLM Model can not be empty!");
        //验证promptID是否为空
        AdePrompt prompt = adePromptMapper.selectAdePromptById(request.getPromptId());
        if(prompt == null) throw new RuntimeException("Prompt can not be null!");

        String systemPrompt = prompt.getSystemPrompt();
        String userPrompt = prompt.getUserPrompt();
        if(StringUtils.isEmpty(systemPrompt)) throw new RuntimeException("System Prompt can not be empty!");
        if(StringUtils.isEmpty(userPrompt)) throw new RuntimeException("User Prompt can not be empty!");

        //获取最终用户提示
        userPrompt = String.format(userPrompt,request.getEmr());
        if(AdePrompt.PROMPT_TYPE_RAG.equals(prompt.getType())){   //如果提示词是RAG模式  那就需要检索知识库
            if(request.getAdekbid()==null) throw new RuntimeException("Knowledge Base ID can not be null!");
            //验证知识库ID是否为空
            AdeKnowledgeBase adeKnowledgeBase = adeKnowledgeBaseMapper.selectAdeKnowledgeBaseById(request.getAdekbid());
            if(adeKnowledgeBase == null) throw new RuntimeException("Knowledge Base can not be null!");
            //验证chunk type是否为空
            if(StringUtils.isEmpty(request.getChunkType())) throw new RuntimeException("Chunk Type can not be empty!");
            if(request.getTopN()>10 || request.getTopN()<1) throw new RuntimeException("Top N can not be greater than 10 or less than 1!");
            if(request.getSimScoreThreshold()>=1 || request.getSimScoreThreshold()<0) throw new RuntimeException("Similarity Score Threshold can not be greater than 1 or less than 0!");

            AdeRetrivedResponse resp = adeKnowledgeEmbedService.search(adeKnowledgeBase.getCollectionName(), adeKnowledgeBase.getEmbedModel(), request.getEmr(), request.getChunkType(), request.getSimScoreThreshold(),request.getTopN(),false, request.isDebug());
            List<AdeRetrivedKnowledge> knowledges = resp.getAdes();
            result.setReferences(knowledges);
            systemPrompt = String.format(systemPrompt, AdeRetrivedKnowledge.formatKnowledge(knowledges));
        }
        result.setPrompt(systemPrompt+userPrompt);
        //调用大语言模型生产
        List<Ade> ades = adeService.recognize(request.getLlmModel(), systemPrompt, userPrompt,true);
        result.setAdes(ades);
        return result;
    }
}
